import torch
import torch.nn as nn
import numpy as np
import random
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from datetime import datetime
import calendar

app = FastAPI()

# CNN+LSTM 预测模型
class CNN_LSTM_Model(nn.Module):
    def __init__(self, input_dim, cnn_out_channels, lstm_hidden_size, output_dim, fixed_length=30):
        super(CNN_LSTM_Model, self).__init__()
        self.fixed_length = fixed_length
        self.conv1d = nn.Conv1d(in_channels=input_dim, out_channels=cnn_out_channels, kernel_size=3, padding=1)
        self.pool = nn.AdaptiveAvgPool1d(output_size=fixed_length)
        self.lstm = nn.LSTM(input_size=cnn_out_channels, hidden_size=lstm_hidden_size, batch_first=True)
        self.fc = nn.Linear(lstm_hidden_size, output_dim)

    def forward(self, x):
        x = [self.conv1d(seq.permute(1, 0).unsqueeze(0)) for seq in x]
        x = [self.pool(seq).squeeze(0).permute(1, 0) for seq in x]
        x = torch.nn.utils.rnn.pad_sequence(x, batch_first=True)
        lstm_out, _ = self.lstm(x)
        output = self.fc(lstm_out[:, -1, :])
        return output

# 训练函数
def train_model(model, data, targets, epochs=5, learning_rate=0.001):
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, torch.tensor(targets, dtype=torch.float32))
        loss.backward()
        optimizer.step()
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

# 滚动预测函数
def rolling_forecast(model, initial_data, predict_start_date):
    predict_start_date = datetime.strptime(predict_start_date, "%Y-%m-%d")
    days_in_month = calendar.monthrange(predict_start_date.year, predict_start_date.month)[1]
    remaining_days = days_in_month - predict_start_date.day + 1

    predictions = []
    current_data = torch.tensor(initial_data, dtype=torch.float32)
    for _ in range(remaining_days):
        output = model([current_data]).item()
        predictions.append(output)
        next_step = torch.tensor([[output] * current_data.shape[1]], dtype=torch.float32)
        current_data = torch.cat((current_data, next_step), dim=0)

    return predictions

# FastAPI 请求格式
class ModelInitRequest(BaseModel):
    input_dim: int
    cnn_out_channels: int
    lstm_hidden_size: int
    output_dim: int
    fixed_length: int

class TrainRequest(BaseModel):
    data: list
    targets: list

class PredictRequest(BaseModel):
    initial_data: list
    predict_start_date: str

class RequestType(BaseModel):
    mode: str  # "train" or "predict"
    train_data: TrainRequest | None = None
    predict_data: PredictRequest | None = None

# FastAPI 端点

@app.post("/initialize_model")
async def initialize_model(model_params: ModelInitRequest):
    global model
    model = CNN_LSTM_Model(
        input_dim=model_params.input_dim,
        cnn_out_channels=model_params.cnn_out_channels,
        lstm_hidden_size=model_params.lstm_hidden_size,
        output_dim=model_params.output_dim,
        fixed_length=model_params.fixed_length
    )
    return {"message": "Model initialized successfully",
            "model_parameters": model_params}


@app.post("/train")
async def train_model_endpoint(request: TrainRequest):
    if model is None:
        raise HTTPException(status_code=400, detail="Model is not initialized")

    data_tensors = [torch.tensor(seq, dtype=torch.float32) for seq in request.data]
    target_tensors = [torch.tensor(target, dtype=torch.float32) for target in request.targets]
    train_model(model, data_tensors, target_tensors)
    return {"message": "Model trained successfully"}


@app.post("/predict")
async def predict(request: PredictRequest):
    if model is None:
        raise HTTPException(status_code=400, detail="Model is not initialized")

    predictions = rolling_forecast(model, request.initial_data, request.predict_start_date)
    return {"predictions": predictions}
